# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Use the distances between atoms to calculate neighbour list
"""

from typing import Tuple
import torch
from torch import Tensor
from torch.nn import Module, Parameter

from ..function import get_integer, get_ms_array, GLOBAL_DEVICE_SPONGE
from ..function import GetDistance


class DistanceNeighbours(Module):
    r"""Neighbour list calculated by distance

    Args:

        cutoff (float):         Cutoff distance.

        num_neighbours (int):   Number of neighbours. If `None` is given, this value will be calculated by
                                the ratio of the number of neighbouring grids to the total number of grids.
                                Default: None

        atom_mask (Tensor):     Tensor of shape `(B, A)`. Data type is bool_.
                                Mask of atoms in the system.
                                Default: None

        exclude_index (Tensor): Tensor of shape `(B, A, Ex)`. Data type is int32.
                                Index of neighbour atoms which could be excluded from the neighbour list.
                                Default: None

        use_pbc (bool):         Whether to use periodic boundary condition. Default: None

        cutoff_scale (float):   Factor to scale the cutoff distance. Default: 1.2

        large_dis (float):      A large number to fill in the distances to the masked neighbouring atoms.
                                Default: 1e4

        cast_fp16 (bool):       If this is set to `True`, the data will be cast to float16 before sort.
                                For use with some devices that only support sorting of float16 data.
                                Default: False

    Supported Platforms:

        ``Ascend`` ``GPU``

    Symbols:

        B:  Number of simulation walker.

        A:  Number of atoms in system.

        N:  Number of neighbour atoms.

        Ex: Maximum number of excluded neighbour atoms.

    """

    def __init__(self,
                 cutoff: float,
                 num_neighbours: int = None,
                 atom_mask: Tensor = None,
                 exclude_index: Tensor = None,
                 use_pbc: bool = None,
                 cutoff_scale: float = 1.2,
                 large_dis: float = 1e4,
                 ):

        super().__init__()

        self.device = GLOBAL_DEVICE_SPONGE()
        self.cutoff = get_ms_array(cutoff, dtype=torch.float32)
        self.cutoff_scale = torch.tensor(cutoff_scale, dtype=torch.float32)
        self.scaled_cutoff = self.cutoff * self.cutoff_scale

        self.num_neighbours = get_integer(num_neighbours)
        if self.num_neighbours is None:
            max_neighbours = torch.tensor(0, dtype=torch.int32)
        else:
            max_neighbours = torch.tensor(self.num_neighbours, dtype=torch.int32)
        self.max_neighbours = get_ms_array(max_neighbours)

        self.large_dis = get_ms_array(large_dis, dtype=torch.float32)

        self.emtpy_atom_shift = 0
        self.atom_mask = None
        self.has_empty_atom = False
        if atom_mask is not None:
            # (B, A)
            self.atom_mask = get_ms_array(atom_mask, dtype=torch.bool)
            if self.atom_mask.ndim == 1:
                self.atom_mask = self.atom_mask.unsqueeze(0)

            self.has_empty_atom = torch.logical_not(self.atom_mask.all())
            if self.has_empty_atom:
                emtpy_atom_mask = torch.logical_not(self.atom_mask)
                # (B, 1, A)
                self.emtpy_atom_shift = emtpy_atom_mask.unsqueeze(-2) * self.large_dis

        self.exclude_index = None
        if exclude_index is not None:
            # (B, A, Ex)
            self.exclude_index = get_ms_array(exclude_index, dtype=torch.int32)
            if self.exclude_index.ndim == 2:
                self.exclude_index = self.exclude_index.unsqueeze(0)

        self.get_distance = GetDistance(use_pbc)

    @staticmethod
    def calc_max_neighbours(distances: Tensor, cutoff: float) -> Tensor:
        """calculate the maximum number of neighbouring atoms"""
        mask = distances.to(cutoff.device) < cutoff
        max_neighbours = torch.count_nonzero(mask.to(torch.float32), -1) - 1
        return torch.max(max_neighbours).to(torch.int32)

    def set_num_neighbours(self,
                           coordinate: Tensor,
                           pbc_box: Tensor = None,
                           scale_factor: float = 1.25
                           ):
        """set maximum number of neighbouring atoms"""
        distances = self.get_distance(coordinate.unsqueeze(-2), coordinate.unsqueeze(-3), pbc_box)
        num_neighbours = self.calc_max_neighbours(distances, self.scaled_cutoff)
        num_neighbours = torch.ceil(get_ms_array(num_neighbours * scale_factor))
        self.num_neighbours = get_integer(torch.min(num_neighbours, torch.tensor(coordinate.shape[-2] - 1)))
        self.max_neighbours.data.copy_(self.num_neighbours)
        return self

    def set_exclude_index(self, exclude_index: Tensor) -> Tensor:
        # (B, A, Ex)
        self.exclude_index = get_ms_array(exclude_index, torch.int32)
        if self.exclude_index.ndim == 2:
            self.exclude_index = self.exclude_index.unsqueeze(0)
        return self.exclude_index

    def print_info(self):
        print(f'[TorchSPONGE] Neighbour list: DistanceNeighbours')
        print(f'[TorchSPONGE]     Cut-off distance: {self.cutoff}')
        print(f'[TorchSPONGE]     Scaled cut-off: {self.scaled_cutoff}')
        print(f'[TorchSPONGE]     Max number of neighbour atoms: {self.num_neighbours}')
        return self

    def check_neighbour_list(self):
        """check the number of neighbouring atoms in neighbour list"""
        if self.num_neighbours is not None and self.max_neighbours > self.num_neighbours:
            raise RuntimeError(f'The max number of neighbour atoms ({self.max_neighbours.numpy()}) is larger than '
                               f'the initial neighbouring number of neighbour list ({self.num_neighbours}!')
        return self

    def forward(self,
                  coordinate: Tensor,
                  pbc_box: Tensor = None,
                  atom_mask: Tensor = None,
                  exclude_index: Tensor = None
                  ) -> Tuple[Tensor, Tensor]:
        r"""Calculate distances and neighbours.

        Args:
            coordinate (Tensor):    Tensor of (B, A, D). Data type is float.
                                    Position coordinates of atoms
            pbc_box (Tensor):       Tensor of (B, D). Data type is bool.
                                    Periodic boundary condition box.
                                    Default: None
            atom_mask (Tensor):     Tensor of (B, A). Data type is bool.
                                    Atomic mask
            exclude_index (Tensor): Tensor of (B, A, Ex). Data type is int.
                                    Index of the atoms that should be excluded from the neighbour list.
                                    Default: None

        Returns:
            distances (Tensor):         Tensor of (B, A, N). Data type is float.
            neighbours (Tensor):        Tensor of (B, A, N). Data type is int.
            neighbour_mask (Tensor):    Tensor of (B, A, N). Data type is bool.

        Symbols:
            B:  Batch size.
            A:  Number of atoms in system.
            N:  Number of neighbour atoms.
            D:  Dimension of position coordinates.
            Ex: Maximum number of excluded neighbour atoms.

        """

        # A
        num_atoms = coordinate.shape[-2]
        # (B, A, A) <- (B, A, 1, 3) - (B, 1, A, 3)
        distances = self.get_distance(coordinate.unsqueeze(-2), coordinate.unsqueeze(-3), pbc_box)

        if atom_mask is None:
            atom_mask = self.atom_mask
            if self.has_empty_atom:
                # (B, A, A) + (B, 1, A)
                distances += self.emtpy_atom_shift
        else:
            if not atom_mask.all():
                emtpy_atom_mask = torch.logical_not(atom_mask)
                # (B, 1, A)
                emtpy_atom_shift = emtpy_atom_mask.unsqueeze(-2) * self.large_dis
                distances = distances + emtpy_atom_shift

        # (B, A)
        if self.num_neighbours is None:
            num_neighbours = num_atoms - 1
        else:
            num_neighbours = self.num_neighbours
            max_neighbours = self.calc_max_neighbours(distances, self.cutoff)
            self.max_neighbours.data.copy_(max_neighbours)

        distances, neighbours = torch.topk(-distances, num_neighbours+1)
        distances = -distances[..., 1:]
        neighbours = neighbours[..., 1:]

        neighbour_mask = distances < self.scaled_cutoff

        if exclude_index is None:
            exclude_index = self.exclude_index
        if exclude_index is not None:
            # (B, A, n, Ex) <- (B, A, n, 1) != (B, A, 1, E)
            exc_mask = neighbours.unsqueeze(-1) != exclude_index.unsqueeze(-2)
            # (B,A,n)
            exc_mask = torch.all(exc_mask, -1)
            neighbour_mask = torch.logical_and(neighbour_mask, exc_mask)

        if atom_mask is not None:
            # (B, A, n) <- (B, A, n) && (B, A, 1)
            neighbour_mask = torch.logical_and(neighbour_mask, atom_mask.unsqueeze(-1))

        # (1, A, 1)
        no_idx = torch.arange(num_atoms, device=self.device).reshape(1, -1, 1)
        # (B, A, n)
        no_idx = torch.broadcast_to(no_idx, neighbours.shape)
        neighbours = torch.where(neighbour_mask, neighbours, no_idx)

        return distances, neighbours, neighbour_mask
